package state

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/google/uuid"
)

// memoryManager implements Manager interface with in-memory storage
type memoryManager struct {
	states map[string]*State
	subs   map[string][]chan *State
	mutex  sync.RWMutex
}

// NewMemoryManager creates a new memory-based state manager
func NewMemoryManager() Manager {
	return &memoryManager{
		states: make(map[string]*State),
		subs:   make(map[string][]chan *State),
	}
}

// Get retrieves a state by ID
func (m *memoryManager) Get(ctx context.Context, id string) (*State, error) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if state, ok := m.states[id]; ok {
		return state, nil
	}
	return nil, fmt.Errorf("state not found: %s", id)
}

// Save saves a state
func (m *memoryManager) Save(ctx context.Context, state *State) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if state.ID == "" {
		state.ID = uuid.NewString()
	}
	state.CreatedAt = time.Now()
	state.UpdatedAt = state.CreatedAt
	state.Version = 1

	m.states[state.ID] = state
	m.notifySubscribers(state.AgentID, state)
	return nil
}

// Update updates a state
func (m *memoryManager) Update(ctx context.Context, state *State) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if existing, ok := m.states[state.ID]; ok {
		state.CreatedAt = existing.CreatedAt
		state.UpdatedAt = time.Now()
		state.Version = existing.Version + 1

		m.states[state.ID] = state
		m.notifySubscribers(state.AgentID, state)
		return nil
	}
	return fmt.Errorf("state not found: %s", state.ID)
}

// Delete deletes a state
func (m *memoryManager) Delete(ctx context.Context, id string) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if state, ok := m.states[id]; ok {
		delete(m.states, id)
		m.notifySubscribers(state.AgentID, &State{
			ID:      id,
			AgentID: state.AgentID,
			Type:    string(EventTypeDeleted),
		})
		return nil
	}
	return fmt.Errorf("state not found: %s", id)
}

// List lists states by agent ID
func (m *memoryManager) List(ctx context.Context, agentID string) ([]*State, error) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	var states []*State
	for _, state := range m.states {
		if state.AgentID == agentID {
			states = append(states, state)
		}
	}
	return states, nil
}

// Watch watches for state changes
func (m *memoryManager) Watch(ctx context.Context, agentID string) (<-chan *State, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	ch := make(chan *State, 100)
	m.subs[agentID] = append(m.subs[agentID], ch)

	go func() {
		<-ctx.Done()
		m.mutex.Lock()
		defer m.mutex.Unlock()

		// Remove subscription when context is done
		subs := m.subs[agentID]
		for i, sub := range subs {
			if sub == ch {
				subs[i] = subs[len(subs)-1]
				subs = subs[:len(subs)-1]
				m.subs[agentID] = subs
				close(ch)
				break
			}
		}
	}()

	return ch, nil
}

// notifySubscribers notifies all subscribers of a state change
func (m *memoryManager) notifySubscribers(agentID string, state *State) {
	if subs, ok := m.subs[agentID]; ok {
		for _, ch := range subs {
			select {
			case ch <- state:
			default:
				// Skip if channel is full
			}
		}
	}
}
